import numpy as np
import scipy.sparse as sp
import pickle
import torch
import torch.sparse as tsp
from torch_geometric.nn.models.autoencoder import negative_sampling
from src.utils import remove_bidirection, to_bidirection

torch.manual_seed(0)  # 设置随机种子后，是每次运行test.py文件的输出结果都一样


def load_data_torch(path, dd_et_list, mono=False):
    """
    :param path: WRITE_DATA_PATH in preprocess_data.py
    :param dd_et_list: a list of int - drug indices  药物指数列表
    :param mono: if consider single drug side effects as drug features 如果考虑单药副作用作为药物特点
    :return: a dict contain: dd-adj list, pp-adj, dp-adj and the feature matrix of drug and protein
    """

    print("loading data")
    # load graph info
    drug_num, protein_num, combo_num, mono_num = 708, 1512, 1, 0
    # 708, 1512, 1, 0

    # ########################################
    # drug-drug 最后的为708*708的药物之间的关系矩阵,能产生什么样的类型
    # ########################################
    dd_list = np.loadtxt("./data/mat_drug_drug.txt")    # 读取的txt文件
    adj = sp.csr_matrix(dd_list)             # 将数据变成csr_matrix类型
    sum_adj = sp.csr_matrix((drug_num, drug_num))
    dd_adj_list = [sp.triu(adj).tocsr()]
    sum_adj += adj

    # ########################################
    # protein-protein 1512*1512  蛋白质之间的稀疏矩阵，之间是否有关系
    # ########################################
    pp_adj = np.loadtxt('./data/mat_protein_protein.txt')
    pp_adj = sp.csr_matrix(pp_adj)

    # ########################################
    # drug-protein    药物与蛋白质之间的稀疏矩阵，之间是否有关系
    # ########################################
    dp_adj = np.loadtxt('./data/mat_drug_protein.txt')
    dp_adj = sp.csr_matrix(dp_adj)

    # # ########################################
    # # remove isolated drugs  移除孤立的药物
    # # ########################################
    # drug_degree = sum_adj.sum(axis=1)   # list大小为，好像是按照1这个轴来计算的
    # isolated_drug = np.where(drug_degree == 0)[0].tolist()  # 要移除的药物的list
    # isolated_num = len(isolated_drug)
    # print("remove ", isolated_num, " isolated drugs: ", isolated_drug)
    # if len(isolated_drug) != 0:
    #     while len(isolated_drug) != 0:
    #         ind = isolated_drug.pop()
    #         # remove from d-d adj   从dd列表中移除
    #         for i in range(len(dd_adj_list)):
    #             dd_adj_list[i] = sp.vstack([dd_adj_list[i][:ind, :],
    #                                         dd_adj_list[i][ind + 1:,
    #                                         :]]).tocsr()
    #             dd_adj_list[i] = sp.hstack([dd_adj_list[i][:, :ind],
    #                                         dd_adj_list[i][:,
    #                                         ind + 1:]]).tocsr()
    #         # remove from d-p adj 从dp列表中移除
    #         dp_adj = sp.vstack([dp_adj[:ind, :], dp_adj[ind + 1:, :]])
    #         # remove from drug additional features   从自己单独产生副作用表移除
    # print('remove finished')
    # ########################################
    # protein feature matrix   蛋白质特征矩阵
    # ########################################
    # protein_feat = sp.identity(protein_num)
    ind = torch.LongTensor([range(protein_num), range(protein_num)])   # [[protein_num 个],[ protein_num 个]] Long类型的张量
    val = torch.FloatTensor([1] * protein_num)                         # 构建一个 [1,1,1...1] 共protein_num个1的张量
    protein_feat = torch.sparse.FloatTensor(ind, val,
                                            torch.Size([protein_num, protein_num]))
    # 标记体为[[protein个],[protein个]] ,value = [1,1,1...1]共protein_num个1

    # ########################################
    # drug feature matrix    药物特征矩阵
    # ########################################
    row = np.array(range(drug_num), dtype=np.long)                # 长度为645的array
    col = np.array(range(drug_num), dtype=np.long)                # 长度为645的array
    mono_num = 0

    ind = torch.LongTensor([row, col])
    val = torch.FloatTensor([1] * len(row))
    drug_feat = torch.sparse.FloatTensor(ind, val,
                                         torch.Size([drug_num, drug_num + mono_num]))

    # return a dict
    data = {'d_feat': drug_feat,
            'p_feat': protein_feat,
            'dd_adj_list': dd_adj_list,
            'dp_adj': dp_adj.tocoo(),
            'pp_adj': pp_adj.tocoo()}

    n_et = 1

    num = [0]
    edge_index_list = []
    edge_type_list = []

    print(n_et, ' polypharmacy side effects')

    for i in range(n_et):
        # pos samples
        adj = dd_adj_list[i].tocoo()
        edge_index_list.append(torch.tensor([adj.row, adj.col], dtype=torch.long))
        edge_type_list.append(torch.tensor([i] * adj.nnz, dtype=torch.long))
        num.append(num[-1] + adj.nnz)

        # if i % 100 == 0:
        #     print(i)

    # data['dd_edge_index'] = t.cat(edge_index_list, 1)
    # data['dd_edge_type'] = t.cat(edge_type_list, 0)
    data['dd_edge_index'] = edge_index_list
    data['dd_edge_type'] = edge_type_list
    data['dd_edge_type_num'] = num
    data['dd_y_pos'] = torch.ones(num[-1])
    data['dd_y_neg'] = torch.zeros(num[-1])

    print('data has been loaded')

    return data


def process_prot_edge(pp_net):
    indices = torch.LongTensor(np.concatenate((pp_net.col.reshape(1, -1),
                                               pp_net.row.reshape(1, -1)),
                                              axis=0))
    indices = remove_bidirection(indices, None)
    n_edge = indices.shape[1]

    rd = np.random.binomial(1, 0.9, n_edge)
    train_mask = rd.nonzero()[0]
    test_mask = (1 - rd).nonzero()[0]

    train_indices = indices[:, train_mask]
    train_indices = to_bidirection(train_indices, None)

    test_indices = indices[:, test_mask]
    test_indices = to_bidirection(test_indices, None)

    return train_indices, test_indices



